/*
 * Copyright (c) 2024 Huawei Device Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */


import { CfgEdge, DfgEdge } from './Edge';
import { Stmt } from "../ArkAnalyzer/src/core/base/Stmt";
import { Local } from '../ArkAnalyzer/src/core/base/Local';
import { AbstractExpr } from '../ArkAnalyzer/src/core/base/Expr';
import { Value } from '../ArkAnalyzer/src';


export class DfgBuilder {
    private cfgNodes: CfgNode[] = [];
    private cfgNodeMap: Map<Stmt, CfgNode> = new Map();
    private dfgEdges: DfgEdge[] = [];
    private changedNodes: Set<CfgNode> = new Set();

    constructor(cfgEdges: CfgEdge[]) {
        if (cfgEdges && cfgEdges.length > 0) {
            this.initNodes(cfgEdges);
            for (let cfgEdge of cfgEdges) {
                this.transNext(cfgEdge);
                this.transFront(cfgEdge);
            }
            this.mapAll();
            this.travel();
            this.buildDfg();
        }
    }

    public getDfgEdges(): DfgEdge[] {
        return this.dfgEdges;
    }

    public initNodes(cfgEdges: CfgEdge[]): void {
        let totalStmt: Set<Stmt> = new Set();
        for (let cfgEdge of cfgEdges) {
            totalStmt.add(cfgEdge.start);
            totalStmt.add(cfgEdge.end);
        }

        for (let stmt of [...totalStmt]) {
            let cfgNode = new CfgNode(stmt);
            this.cfgNodes.push(cfgNode);
            this.cfgNodeMap.set(stmt, cfgNode);
        }
    }

    public mapAll(): void {
        for (let cfgNode of this.cfgNodes) {
            for (let next of cfgNode.getNext()) {
                let tmp = this.cfgNodeMap.get(next);
                if (tmp) {
                    cfgNode.addNextNodes(tmp);
                }
            }
        }
    }

    public transNext(cfgEdge: CfgEdge): void {
        for (let cfgNode of this.cfgNodes) {
            if (cfgNode.getSelf() === cfgEdge.start) {
                cfgNode.addNext(cfgEdge.end);
                return;
            }
        }
    }

    public transFront(cfgEdge: CfgEdge): void {
        for (let cfgNode of this.cfgNodes) {
            if (cfgNode.getSelf() === cfgEdge.end) {
                cfgNode.addFront(cfgEdge.start);
                return;
            }
        }
    }

    public resetAll(): void {
        for (let cfgNode of this.cfgNodes) {
            cfgNode.setVisited(false);
            cfgNode.setFrontSize();
        }
        this.changedNodes = new Set();
    }

    public getOne(): CfgNode | null {
        let minOut: CfgNode | null = null;
        for (let cfgNode of this.cfgNodes) {
            if (cfgNode.getVisited()) {
                continue;
            }
            else if (cfgNode.getFrontSize() === 0) {
                return cfgNode;
            }
            else if (minOut === null || cfgNode.getFrontSize() < minOut.getFrontSize()) {
                minOut = cfgNode;
            }
        }
        return minOut;
    }

    public defStmtChanged(before: Set<Stmt>, after: Set<Stmt>): boolean {
        if (before.size !== after.size) {
            return true;
        }
        return false;
    }

    public visit(cfgNode?: CfgNode): boolean {
        if (cfgNode === undefined) {
            let tmp: CfgNode | null = this.getOne();
            if (tmp === null) {
                return false;
            }
            else {
                cfgNode = tmp;
            }
        }
        if (cfgNode.getVisited()) {
            return true;
        }
        cfgNode.setVisited(true);

        for (let nextNode of cfgNode.getNextNodes()) {
            if (nextNode instanceof CfgNode) {
                let tmp: Set<Stmt> = nextNode.getDefStmt();
                nextNode.setDefStmt(new Set([...nextNode.getDefStmt(), ...cfgNode.getDefStmt(), cfgNode.getSelf()]));
                nextNode.declineFrontSize();
                if (this.defStmtChanged(tmp, nextNode.getDefStmt())) {
                    this.changedNodes.add(nextNode);
                }
                if (nextNode.getFrontSize() === 0) {
                    this.visit(nextNode);
                }
            }
        }
        return true;
    }

    public travel(): void {
        this.changedNodes = new Set(this.cfgNodes);
        while (this.changedNodes.size !== 0) {
            this.resetAll();
            while (this.visit());
        }
    }
    public addUse(use: Value, uses: Set<Local>): void {
        if (use instanceof Local) {
            uses.add(use);
        }
        else if (use instanceof AbstractExpr) {
            let expr = use as AbstractExpr;
            for (let simple of expr.getUses()) {
                if (simple instanceof Local) {
                    uses.add(simple);
                }
            }
        }
    }
    public addDefStmt(cfgNode: CfgNode, uses: Set<Local>): void {
        for (let defStmt of cfgNode.getDefStmt()) {
            let def = defStmt.getDef();
            if (def) {
                let local = def as Local;
                if (uses.has(local)) {
                    this.dfgEdges.push(new DfgEdge(cfgNode.getSelf(), defStmt));
                }
            }
        }
    }

    public buildDfg(): void {
        for (let cfgNode of this.cfgNodes) {
            for (let use of cfgNode.getSelf().getUses()) {
                let uses: Set<Local> = new Set();
                this.addUse(use, uses);
                this.addDefStmt(cfgNode, uses);
            }
        }
    }
}

class CfgNode {
    private self: Stmt;
    private next: Stmt[] = [];
    private front: Stmt[] = [];
    private visited: boolean = false;
    private frontSize: number = 0;
    private defStmt: Set<Stmt> = new Set();
    private nextNodes: CfgNode[] = [];

    constructor(self: Stmt) {
        this.self = self;
    }

    public addNext(next: Stmt): void {
        this.next.push(next);
    }

    public addFront(front: Stmt): void {
        this.front.push(front);
    }

    public getNext(): Stmt[] {
        return this.next;
    }

    public addNextNodes(cfgNode: CfgNode) {
        this.nextNodes.push(cfgNode);
    }

    public getSelf(): Stmt {
        return this.self;
    }

    public setVisited(visited: boolean) {
        this.visited = visited;
    }

    public setFrontSize() {
        this.frontSize = this.front.length;
    }

    public getVisited(): boolean {
        return this.visited;
    }

    public getFrontSize(): number {
        return this.frontSize;
    }

    public getNextNodes(): CfgNode[] {
        return this.nextNodes;
    }

    public getDefStmt(): Set<Stmt> {
        return this.defStmt;
    }

    public setDefStmt(stmts: Set<Stmt>) {
        this.defStmt = stmts;
    }

    public declineFrontSize() {
        this.frontSize--;
    }
}
